package net.pragyah.scalby

import junit.framework.TestCase
import junit.framework.Assert._
 
import net.pragyah.scalby._


class TestProbabilityDistribution extends TestCase{
  
	val x1 = new Variable[int]("X1",Set[int]()+1+2,None) 
  	val x2 = new Variable[int]("X2",Set[int]()+1+2,None) 
  	val x3 = new Variable[int]("X3",Set[int]()+1+2,None) 
  	val x4 = new Variable[int]("X4",Set[int]()+1+2,None) 
  	val x5 = new Variable[int]("X5",Set[int]()+1+2,None) 

  def testCreationAndQuery {
    val pd = ProbabilityDistribution[int]()
    
    pd(x2,1,Map((x1,1)),0.8f)
    assert(pd(x2,1,Map((x1,1))) != None)
    assertEquals(0.8f,pd(x2,1,Map((x1,1))).get)

    assert(pd(x2,1,Map((x1,2))) == None)

	pd(x2,1,Map((x1,1)) +(((x2,2))) ,0.5f)
    assert(pd(x2,1,Map((x1,1)) +(((x2,2)))) != None)
    assertEquals(0.5f,pd(x2,1,Map((x1,1)) +(((x2,2)))).get)

    assert(pd(x2,1,Map((x1,2)) +(((x2,2)))) == None)
    
  }
  
  
  def testDependencies {
    val pd = ProbabilityDistribution[int]()
    
    pd(x5,1,Map((x2,1))+((x3,1)),0.1f)
    pd(x5,1,Map((x2,1))+((x3,2)),0.2f)
    pd(x5,1,Map((x2,2))+((x3,1)),0.3f)
    pd(x5,1,Map((x2,2))+((x3,2)),0.4f)
    pd(x5,2,Map((x2,1))+((x3,1)),0.9f)
    pd(x5,2,Map((x2,1))+((x3,2)),0.8f)
    pd(x5,2,Map((x2,2))+((x3,1)),0.7f)
    pd(x5,2,Map((x2,2))+((x3,2)),0.6f)
    
    
    val dep = pd.getDependencies(x5,1)
    assertEquals(4,dep.size)
    assert(dep.contains(Map((x2,1))+((x3,1))))
    assert(dep.contains(Map((x2,1))+((x3,2))))
    assert(dep.contains(Map((x2,2))+((x3,1))))
    assert(dep.contains(Map((x2,2))+((x3,2))))
    
    assertEquals(0.1f,dep(Map((x2,1))+((x3,1))))
    assertEquals(0.2f,dep(Map((x2,1))+((x3,2))))
    assertEquals(0.3f,dep(Map((x2,2))+((x3,1))))
    assertEquals(0.4f,dep(Map((x2,2))+((x3,2))))
 
    
    val dep2 = pd.getDependencies(x5,2)
    assertEquals(4,dep2.size)
    assert(dep2.contains(Map((x2,1))+((x3,1))))
    assert(dep2.contains(Map((x2,1))+((x3,2))))
    assert(dep2.contains(Map((x2,2))+((x3,1))))
    assert(dep2.contains(Map((x2,2))+((x3,2))))
    
    assertEquals(0.9f,dep2(Map((x2,1))+((x3,1))))
    assertEquals(0.8f,dep2(Map((x2,1))+((x3,2))))
    assertEquals(0.7f,dep2(Map((x2,2))+((x3,1))))
    assertEquals(0.6f,dep2(Map((x2,2))+((x3,2))))
  }
  
  
  def testOtherDependencies {
    val pd = ProbabilityDistribution[int]()
    
    pd(x5,1,Map((x2,1))+((x3,1)),0.1f)
    pd(x5,1,Map((x2,1))+((x3,2)),0.2f)
    pd(x5,1,Map((x2,2))+((x3,1)),0.3f)
    pd(x5,1,Map((x2,2))+((x3,2)),0.4f)
    pd(x5,2,Map((x2,1))+((x3,1)),0.9f)
    pd(x5,2,Map((x2,1))+((x3,2)),0.8f)
    pd(x5,2,Map((x2,2))+((x3,1)),0.7f)
    pd(x5,2,Map((x2,2))+((x3,2)),0.6f)
    
	var dep = pd.getOtherDependencies(x5,2,x3,2)
    assertEquals(2,dep.size)
    assert(dep.contains(Map((x2,1))))
    assert(dep.contains(Map((x2,2))))
    assertEquals(0.8f,dep(Map((x2,1))))
    assertEquals(0.6f,dep(Map((x2,2))))
    
    
	dep = pd.getOtherDependencies(x5,1,x2,2)
    assertEquals(2,dep.size)
    assert(dep.contains(Map((x3,1))))
    assert(dep.contains(Map((x3,2))))
    assertEquals(0.3f,dep(Map((x3,1))))
    assertEquals(0.4f,dep(Map((x3,2))))
    
    
  }
  
  def testGetConditionalForAll{
    val pd = ProbabilityDistribution[int]()
    
    pd(x5,1,Map((x2,1))+((x3,1)),0.1f)
    pd(x5,1,Map((x2,1))+((x3,2)),0.2f)
    pd(x5,1,Map((x2,2))+((x3,1)),0.3f)
    pd(x5,1,Map((x2,2))+((x3,2)),0.4f)
    pd(x5,2,Map((x2,1))+((x3,1)),0.9f)
    pd(x5,2,Map((x2,1))+((x3,2)),0.8f)
    pd(x5,2,Map((x2,2))+((x3,1)),0.7f)
    pd(x5,2,Map((x2,2))+((x3,2)),0.6f)
    
    var aA = Map[Variable[int],int]() + (x1 -> 2) + (x2 -> 1) + (x3 -> 2) + (x4 -> 2) 
	var prob = pd(x5,1,aA)
    assert(prob != None)
    assertEquals(0.2f,prob.get)

    
    aA = Map[Variable[int],int]() + (x1 -> 2) + (x2 -> 1) + (x4 -> 2) 
    prob = pd(x5,1,aA)
    assertEquals(None,prob)
    
    
    aA = Map[Variable[int],int]() + (x1 -> 2) + (x4 -> 2) 
    prob = pd(x5,1,aA)
    assertEquals(None,prob)
    
    
  }

}
